package forward_split

import (
	"flag"
	"fmt"
	"github.com/panjf2000/ants/v2"
	log "github.com/sirupsen/logrus"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"
)

type Controler struct {
	Running        bool
	Client_conn    net.Conn
	Client_data    chan []byte
	Client_rw_lock sync.RWMutex

	Remote_addr      string
	Remote_addr_list []string
	Remote_conn      net.Conn
	Remote_data      chan []byte
	Remote_rw_lock   sync.RWMutex
}

var (
	channel_size = 10240
	g_pool_size  int
)

func connectToAddr(remote string) (net.Conn, error) {
	return net.Dial("tcp", remote)
}

func (c *Controler) reconnectToRemote() error {
	for times, interval := range []byte{1, 1, 2, 2, 4, 4, 8, 8, 16, 16, 32, 32} {
		func() {
			c.Remote_rw_lock.Lock()
			defer c.Remote_rw_lock.Unlock()

			addr := c.Remote_addr_list[rand.Int()%len(c.Remote_addr_list)]
			c.Remote_addr = addr
		}()

		remote_conn, err := connectToAddr(c.Remote_addr)
		if err == nil {
			func() {
				c.Remote_rw_lock.Lock()
				defer c.Remote_rw_lock.Unlock()

				c.Remote_conn = remote_conn
			}()

			return nil
		}

		log.Trace("reconnectToRemote error: ", err, "and try times: ", times+1)
		time.Sleep(time.Second * time.Duration(interval))
	}

	return fmt.Errorf("error connect to remote: %s", c.Remote_addr)
}

func NewControler(client net.Conn, remote []string) *Controler {
	var (
		con net.Conn
		err error
	)

	if len(remote) <= 0 {
		return nil
	}

	if con, err = connectToAddr(remote[0]); err != nil {
		log.Errorf("connectToAddr %s error\n", remote)
	}
	return &Controler{
		Running: false,

		Client_conn:    client,
		Client_data:    make(chan []byte, channel_size),
		Client_rw_lock: sync.RWMutex{},

		Remote_addr:      remote[0],
		Remote_addr_list: remote,
		Remote_conn:      con,
		Remote_data:      make(chan []byte, channel_size),
		Remote_rw_lock:   sync.RWMutex{},
	}
}

func (c *Controler) Close() {
	c.Running = false

	close(c.Client_data)
	close(c.Remote_data)
}

func (c *Controler) recClientInfo() {
	buffers := make([]byte, 1024)

	for c.Running {
		var err error
		rc, err := c.Client_conn.Read(buffers)
		if err != nil {
			log.Tracef("c.client_conn.Read error: ", err)

			c.Close()
			return
		} else {
			if rc > 0 {
				func() {
					defer func() {
						if err := recover(); err != nil {
							log.Error(err)
							return
						}
					}()
					c.Client_data <- buffers[:rc]
					log.Tracef("c.client_conn.Read: ", string(buffers[:rc]))
				}()

			}
		}

	}
}

func (c *Controler) sendClientInfo() {
	for c.Running {
		for data := range c.Client_data {
			if c.Remote_conn == nil {
				err := c.reconnectToRemote()
				if err != nil {
					log.Error(err)
					return
				}
			}

			count, err := c.Remote_conn.Write(data)
			if err != nil {
				log.Tracef("c.remote_conn.Write error: ", err)
				err := c.reconnectToRemote()
				if err != nil {
					log.Error(err)
					return
				}

				count, _ = c.Remote_conn.Write(data)
			} else {
				log.Tracef("c.remote_conn.Write success and count: ", count)
			}
		}

	}
}

func (c *Controler) recRemoteInfo() {
	buffers := make([]byte, 1024)

	for c.Running {
		var rc int
		var err error
		func() {
			if c.Remote_conn == nil {
				err := c.reconnectToRemote()
				if err != nil {
					log.Error(err)
					return
				}
			}
		}()

		func() {
			c.Remote_rw_lock.RLock()
			defer c.Remote_rw_lock.RUnlock()
			rc, err = c.Remote_conn.Read(buffers)
		}()

		if err != nil {
			log.Error("c.remote_conn.Read error: ", err)

			err := c.reconnectToRemote()
			if err != nil {
				log.Error(err)
				return
			}

			func() {
				c.Remote_rw_lock.RLock()
				defer c.Remote_rw_lock.RUnlock()
				rc, _ = c.Remote_conn.Read(buffers)
			}()
		} else {
			if rc > 0 {
				func() {
					defer func() {
						if err := recover(); err != nil {
							log.Error(err)
							return
						}
					}()

					c.Remote_data <- buffers[:rc]
					log.Trace("c.remote_conn.Read: ", string(buffers[:rc]))
				}()

			}
		}

	}
}

func (c *Controler) sendRemoteInfo() {
	for c.Running {
		for data := range c.Remote_data {
			count, err := c.Client_conn.Write(data)
			if err != nil {
				log.Error("c.client_conn.Write error: ", err)
			} else {
				log.Trace("c.client_conn.Write success and count: ", count)
			}
		}

	}
}

func (c *Controler) Run() {
	c.Running = true

	go c.recClientInfo()
	go c.recRemoteInfo()

	go c.sendClientInfo()
	go c.sendRemoteInfo()
}

func Run() error {
	var local_port, remote_addr_list string
	var pool_size int

	flag.StringVar(&local_port, "local_port", "50050", "--local_port 50050")
	flag.StringVar(&remote_addr_list, "server", ":50051,:50052", "--server :50051,:50052")
	flag.IntVar(&pool_size, "pool_size", 100*100, "--pool_size 100*100")

	g_pool_size = pool_size
	flag.Parse()

	remote_addr := strings.Split(remote_addr_list, ",")
	log.Tracef("local_port: %s remote_addr: %s\n", local_port, remote_addr)

	local_addr := fmt.Sprintf("0.0.0.0:%s", local_port)
	log.Trace("start to listen local addr: ", local_addr)
	server, err := net.Listen("tcp", local_addr)
	if err != nil {
		log.Errorf("listen error: %s", err.Error())
		return err
	}
	log.Trace("listen success: ", local_addr)

	pool, err := ants.NewPool(g_pool_size)
	if err != nil {
		return err
	}

	for {
		conn, err := server.Accept()
		if err != nil {
			log.Errorf("accept error: %s", err.Error())
			continue
		}
		log.Tracef("accept from: %s\n", conn.RemoteAddr())

		_ = conn.(*net.TCPConn).SetKeepAlive(true)
		client := NewControler(conn, remote_addr)
		_ = pool.Submit(client.Run)
	}
}
